package com.twobu.twobuinsight.aspect;

import com.twobu.twobuinsight.annotation.RateLimit;

import com.twobu.twobuinsight.exception.BusinessException;
import com.twobu.twobuinsight.security.UserContext;
import com.twobu.twobuinsight.security.UserContextHolder;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 限流切面
 */
@Aspect
@Component
@Slf4j
public class RateLimitAspect {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    private static final String RATE_LIMIT_KEY_PREFIX = "rate:limit:";

    /**
     * 限流Lua脚本
     * 使用Redis的INCR和EXPIRE命令实现计数器限流算法
     */
    private static final String RATE_LIMIT_LUA_SCRIPT =
            "local key = KEYS[1] " +
            "local count = tonumber(ARGV[1]) " +
            "local timeValue = tonumber(ARGV[2]) " +
            "local current = redis.call('incr', key) " +
            "if current == 1 then " +
            "    redis.call('expire', key, timeValue) " +
            "end " +
            "if current > count then " +
            "    return 0 " +
            "end " +
            "return 1";

    /**
     * 环绕通知处理限流逻辑
     */
    @Around("@annotation(com.twobu.twobuinsight.annotation.RateLimit)")
    public Object rateLimit(ProceedingJoinPoint joinPoint) throws Throwable {
        // 获取方法签名
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        // 获取限流注解
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        if (rateLimit == null) {
            return joinPoint.proceed();
        }

        // 获取当前用户上下文
        UserContext userContext = UserContextHolder.getUserContext();
        if (userContext == null || userContext.getUserId() == null) {
            log.warn("无法获取用户上下文，跳过限流检查");
            return joinPoint.proceed();
        }

        // 构建限流Key: 前缀 + 用户ID + 类名 + 方法名
        String className = method.getDeclaringClass().getName();
        String methodName = method.getName();
        String key = RATE_LIMIT_KEY_PREFIX + userContext.getUserId() + ":" + className + ":" + methodName;

        // 获取限流参数
        int count = rateLimit.count();
        int timeValue = rateLimit.timeValue();
        TimeUnit timeUnit = rateLimit.timeUnit();

        // 转换时间单位为秒
        long seconds = timeUnit.toSeconds(timeValue);

        // 执行Lua脚本进行限流检查
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptText(RATE_LIMIT_LUA_SCRIPT);
        redisScript.setResultType(Long.class);

        List<String> keys = Collections.singletonList(key);
        Long result = stringRedisTemplate.execute(redisScript, keys, String.valueOf(count), String.valueOf(seconds));

        // 如果结果为0，表示超过限流阈值
        if (result != null && result == 0) {
            log.warn("用户[{}]调用方法[{}.{}]超过限流阈值: {} 次/{} {}",
                    userContext.getUserId(), className, methodName, count, timeValue, timeUnit.name().toLowerCase());
            throw new BusinessException(rateLimit.description());
        }

        // 正常执行方法
        return joinPoint.proceed();
    }
}
